//
// Created by ZhaoXiaoFei on 2022/7/17.
//

#include "feature_manager.hpp"


int FeatureManager::getFeatureCount(){
    int cnt = 0;
    for(auto& it : allFeatures){
        it.used_num = it.features_per_id.size();
        if(it.used_num >= 2 && it.start_frame < 10 - 2){
            cnt++;
        }
    }
    return cnt;
}

Eigen::VectorXd FeatureManager::getDepthVector(){
    Eigen::VectorXd res(getFeatureCount());
    int featureIndex = -1;
    for(auto &it : allFeatures){
        it.used_num = it.features_per_id.size();
        if(!(it.used_num >= 2 && it.start_frame < 10 - 2)){
            continue;
        }
        res(++featureIndex) = 1.0 / it.estimated_depth;
    }
    return res;
}

void FeatureManager::setDepthVector(const Eigen::VectorXd &x) {
    int featureIndex = -1;
    for(auto &it : allFeatures){
        it.used_num = it.features_per_id.size();
        if(!(it.used_num >= 2 && it.start_frame < 10 - 2)){
            continue;
        }
        it.estimated_depth = 1.0 / x(++featureIndex);
        if(it.estimated_depth < 0){
            it.solve_flag = 2;
        }
        else{
            it.solve_flag = 1;
        }
    }
}

void FeatureManager::clearDepth(Eigen::VectorXd& depthVec){
    int featureIndex = -1;
    for(auto &it : allFeatures){
        it.used_num = it.features_per_id.size();
        if(!(it.used_num >= 2 && it.start_frame < 10 - 2)){
            continue;
        }
        it.estimated_depth = 1.0 / depthVec(++featureIndex);
    }
}
void FeatureManager::triangulation(Eigen::Vector3d Ps[], Eigen::Vector3d& tbc, Eigen::Matrix3d& rbc){
//    LOG(INFO) << "tbc: " << tbc << std::endl;
//    LOG(INFO) << "rbc: " << rbc << std::endl;
    for(auto& it_per_ids : allFeatures){
        it_per_ids.used_num = it_per_ids.features_per_id.size();
        if(!(it_per_ids.used_num >= 2 && it_per_ids.start_frame < 10 - 2)){
            continue;
        }
        if(it_per_ids.estimated_depth > 0){
            continue;
        }
        int imu_i = it_per_ids.start_frame;
        int imu_j = imu_i - 1;
        Eigen::MatrixXd A(2 * it_per_ids.features_per_id.size(), 4);
        Eigen::Vector3d T0 = Ps[imu_i] + Rs[imu_i] * tbc;
        Eigen::Matrix3d R0 = Rs[imu_i] * rbc;
        Eigen::Matrix<double, 3, 4> P;
        int index = 0;
        for(auto& it_per_frame : it_per_ids.features_per_id){
            imu_j++;
            Eigen::Vector3d T1 = Ps[imu_j] + Rs[imu_j] * tbc;
            Eigen::Matrix3d R1 = Rs[imu_j] * rbc;
            Eigen::Vector3d relativeT = R0.transpose() * (T1 - T0);
            Eigen::Matrix3d relativeR = R0.transpose() * R1;
//            LOG(INFO) << "relativeR: " << Eigen::Quaterniond(relativeR) << std::endl;
//            LOG(INFO) << "relativeT: " << relativeT.transpose() << std::endl;
            P.block<3, 3>(0, 0) = relativeR.transpose();
            P.block<3, 1>(0, 3) = -relativeR.transpose() * relativeT;
            Eigen::Vector3d pointW = it_per_frame.point.normalized();
//            LOG(INFO) << "pointW: " << pointW << std::endl;
            A.row(index++) = P.row(2) * pointW[0] - P.row(0) * pointW[2];
            A.row(index++) = P.row(2) * pointW[1] - P.row(1) * pointW[2];
            if(imu_i == imu_j){
                continue;
            }
        }

        Eigen::Vector4d x = Eigen::JacobiSVD<Eigen::MatrixXd>(A, Eigen::ComputeFullU | Eigen::ComputeFullV).matrixV().rightCols<1>();
//        LOG(INFO) << "x: " << x.transpose() << std::endl;
        double D = x(2) / x(3);
        it_per_ids.estimated_depth = D;
        if(it_per_ids.estimated_depth < 0.1){
            it_per_ids.estimated_depth = 5.0;
        }
//        LOG(INFO) << "estimated_depth: " << it_per_ids.estimated_depth << std::endl;

    }
}

double FeatureManager::computeParallax(FeaturesPerFeatureId& it, int frame_count){
    const FeaturePerFrame& frame_i = it.features_per_id[frame_count - it.start_frame - 2];
    const FeaturePerFrame& frame_j = it.features_per_id[frame_count - it.start_frame - 1];

    double res = 0.0;
    Eigen::Vector3d p_i = frame_i.point;
    double u_i = p_i.x() / p_i.z();
    double v_i = p_i.y() / p_i.z();

    Eigen::Vector3d p_j = frame_j.point;
    double u_j = p_j.x() / p_j.z();
    double v_j = p_j.y() / p_j.z();

    double dx = u_j - u_i;
    double dy = v_j - v_i;

    res = std::max(res, std::sqrt(dx * dx + dy * dy));

    return res;
    //check the second last frame is keyframe or not
    //parallax betwwen seconde last frame and third last frame
}

bool FeatureManager::addFeatureAndCheckParallax(int frame_count, const std::map<int, std::vector<std::pair<int, Eigen::Matrix<double, 7, 1>>>>& image){
    int track_count = 0;
    for(auto& i : image){
        int featureId = i.first;
        //i.second.size() === 1;
        FeaturePerFrame featurePerFrame(i.second[0].second);
        auto it = std::find_if(allFeatures.begin(), allFeatures.end(), [featureId](const FeaturesPerFeatureId& it){
            return it.featurId == featureId;
        });
        if(it == allFeatures.end()){
            allFeatures.push_back(FeaturesPerFeatureId(featureId, frame_count));
            allFeatures.back().features_per_id.push_back(featurePerFrame);
        }
        else if(it->featurId == featureId){
            it->features_per_id.push_back(featurePerFrame);
            track_count++;
        }
    }
    if(frame_count < 2 || track_count < 20){
        return true;
    }
    double parallax_sum = 0.0;
    int parallax_num = 0;
    for(auto& it : allFeatures){
        if(it.start_frame <= frame_count - 2 &&
           it.start_frame + int(it.features_per_id.size()) - 1 >= frame_count - 1){
            parallax_num++;
            parallax_sum += computeParallax(it, frame_count);
        }
    }
    if(parallax_num == 0){
        return true;
    }
    else{
        LOG(INFO) << "parallax_sum: " << parallax_sum << " parallax_num: " << parallax_num << std::endl;
        LOG(INFO) << "current parallax: " << parallax_sum / parallax_num * 460.0 << std::endl;
        return (parallax_sum / parallax_num) >= (10.0 / 460.0);
    }
}

void FeatureManager::removeOld()
{

    for(auto it = allFeatures.begin(), it_next = allFeatures.begin(); it != allFeatures.end(); it = it_next){

        it_next++;

        if(it->start_frame != 0){
            it->start_frame--;
        }else{
            it->features_per_id.erase(it->features_per_id.begin());
            if(it->features_per_id.size() == 0){
                allFeatures.erase(it);
            }
        }
    }
}

void FeatureManager::removeFailures() {
    for (auto it = allFeatures.begin(), it_next = allFeatures.begin();
         it != allFeatures.end(); it = it_next)
    {
        it_next++;
        if (it->solve_flag == 2)
            allFeatures.erase(it);
    }
}

void FeatureManager::removeOldAndShiftDepth(Eigen::Matrix3d &R1, Eigen::Vector3d &T1, Eigen::Matrix3d &R2,
                                            Eigen::Vector3d &T2) {

    for(auto it = allFeatures.begin(), it_next = allFeatures.begin(); it != allFeatures.end(); it = it_next){

        it_next++;

        if(it->start_frame != 0){
            it->start_frame--;
        }else{
            Eigen::Vector3d uv_point = it->features_per_id[0].point;
            it->features_per_id.erase(it->features_per_id.begin());
            if(it->features_per_id.size() < 2){
                allFeatures.erase(it);
                continue;
            }
            else{
                Eigen::Vector3d cam_point_i = uv_point * it->estimated_depth;
                Eigen::Vector3d w_point = R1 * cam_point_i + T1;
                Eigen::Vector3d cam_point_j = R2.transpose() * (w_point - T2);
                double depth_j = cam_point_j(2);
                if(depth_j > 0){
                    it->estimated_depth = depth_j;
                }
                else{
                    it->estimated_depth = 5.0;
                }
            }
        }
    }
}

void FeatureManager::clearState() {
    allFeatures.clear();
}

void FeatureManager::removeSecondNew(int frame_count){

    for(auto it = allFeatures.begin(), it_next = allFeatures.begin(); it != allFeatures.end(); it = it_next){
        it_next++;
        if(it->start_frame == frame_count){
            it->start_frame--;
        }
        else{
            if(it->endFrame() < frame_count - 1){
                continue;
            }
            int j = 10 - 1 - it->start_frame;
            it->features_per_id.erase(it->features_per_id.begin() + j);
            if(it->features_per_id.size() == 0){
                allFeatures.erase(it);
            }
        }
    }
}

std::vector<std::pair<Eigen::Vector3d, Eigen::Vector3d>>
FeatureManager::getCorresponding(int frameIdL, int frameIdR) {
    std::vector<std::pair<Eigen::Vector3d, Eigen::Vector3d>> res;
    for(auto& it : allFeatures){
        if(it.start_frame <= frameIdL && it.endFrame() >= frameIdR){
            Eigen::Vector3d leftP = it.features_per_id[frameIdL - it.start_frame].point;
            Eigen::Vector3d rightP = it.features_per_id[frameIdR - it.start_frame].point;
            res.push_back(std::make_pair(leftP, rightP));
        }
    }
    return res;
}